package chat

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"net/http"

	"github.com/gin-gonic/gin"
)

// OllamaAPIURL 是 Ollama 的 API 地址
const OllamaAPIURL = "http://localhost:11434/api/generate"

// ChatRequest 定义了客户端发送给后端的请求结构体
type ChatRequest struct {
	Prompt string `json:"prompt"`
}

// OllamaResponse 定义了 Ollama 返回的响应结构体
type OllamaResponse struct {
	Response string `json:"response"`
	Done     bool   `json:"done"`
}

// OllamaHandleChat 处理客户端的聊天请求，并与 Ollama 进行交互
func OllamaHandleChat(c *gin.Context) {
	// 初始化请求结构体
	var req ChatRequest
	// 从请求体中解析 JSON 数据到请求结构体
	if err := c.ShouldBindJSON(&req); err != nil {
		// 若解析失败，返回 400 错误
		c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
		return
	}

	// 构建发送给 Ollama 的请求体
	ollamaReqBody := map[string]interface{}{
		"model":  "deepseek-r1:1.5b", // 可根据实际使用的模型修改
		"prompt": req.Prompt,
	}
	// 将请求体转换为 JSON 字节切片
	ollamaReqJSON, err := json.Marshal(ollamaReqBody)
	if err != nil {
		// 若转换失败，返回 500 错误
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}

	// 创建一个 POST 请求到 Ollama 的 API
	resp, err := http.Post(OllamaAPIURL, "application/json", bytes.NewReader(ollamaReqJSON))
	if err != nil {
		// 若请求失败，返回 500 错误
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}
	// 确保在函数结束时关闭响应体
	defer resp.Body.Close()

	// 读取 Ollama 的响应体
	body, err := io.ReadAll(resp.Body)
	if err != nil {
		// 若读取失败，返回 500 错误
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}

	// 打印接收到的响应内容，用于调试
	fmt.Printf("Received response from Ollama: %s\n", string(body))

	var combinedResponse string
	decoder := json.NewDecoder(bytes.NewReader(body))
	for {
		var ollamaResp OllamaResponse
		if err := decoder.Decode(&ollamaResp); err != nil {
			if err == io.EOF {
				break
			}
			// 若解析失败，返回 500 错误
			c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
			return
		}
		combinedResponse += ollamaResp.Response
		if ollamaResp.Done {
			break
		}
	}

	// 返回成功响应给客户端
	c.JSON(http.StatusOK, gin.H{"response": combinedResponse})
}
